#pragma OPENCL EXTENSION cl_khr_fp64 : enable

#include <params.hcl>

//Default wave function
#if defined(DEFAULT_WAVE_FUNCTION)
float computeLogPsi(float r)
{
  const float curveFit = 0.63837943620287906921f;
  float ri;

  ri = 1.0f / (r - p4);
  float riPow2 = ri * ri;

  return -1.0f * ((p0 * riPow2 * ri + p1) * riPow2 + p2 * log(r)) + curveFit;
}

float computeLogPsiPrime(float r)
{
  float rInv = 1.0 / (r - p4);
  float rInvPow3 = rInv * rInv * rInv;

  return (5.0 * p0 * rInvPow3 + 2.0 * p1) * rInvPow3 - p2 * rInv - p3;
}

float computeLogPsiDoublePrime(float r)
{
  float rInv = 1.0 / (r - p4);
  float rInvPow2 = rInv * rInv;
  float rInvPow3 = rInvPow2 * rInv;

  return (-1.0 * (30.0 * p0 * rInvPow3 + 6.0 * p1) * rInvPow2 - p2) * rInvPow2;
}

//McMillan Wave functions and derrivatives
#elif defined(MCMILLAN_WAVE_FUNCTION)
float computeLogPsi(float r)
{
  return c00 * pow(r, c10) + c20 * r * r;
}

float computeLogPsiPrime(float r)
{
  return c01 * pow(r, c11) + c21 * r;
}

float computeLogPsiDoublePrime(float r)
{
  return c02 * pow(r, c12) + c22;
}

#else
#error "kernel.cl:Invalid wave function."
#endif

float computePotential(float r)
{
	float x = r/rm;
	float xInverse = 1.0f / x;
	float xInversePow2 = xInverse * xInverse;
	float xInversePow6 = xInversePow2 * xInversePow2 * xInversePow2;
	float potential1 = A * exp(x * (beta * x - alphaPotential));
	float potential2 = (c6_pot + 
                      (c8_pot + c9_pot*xInverse + c10_pot * xInversePow2) * 
											xInversePow2) * xInversePow6;
	float potential3 = 1.0f;

	if(x < D)
	{
		float tmp = (D / x) - 1.0f;
		potential3 = exp(-1.0f * tmp * tmp);
	}

	return epsilon * (potential1 - (potential2 * potential3));

}

#ifdef GPU_SPRNG

//sprng crap
__kernel void init_sprng(__global rngen* generators,
                         unsigned int walkerStart,
                         unsigned int numWalkers,
                         __global CLSprngErr* err)
{
  unsigned int thisWalker = get_global_id(0);
  CLSprngErr error = (int)numWalkers;

  if(thisWalker < numWalkers)
  {
    /*
    error = init_rng(walkerStart + thisWalker, 
                     numWalkers, 
                     &generators[thisWalker], 
                     seed, 
                     0);
                     */
    if(error != kCLSprngErrNoError)
    {
      *err = error;
    }
  }
}

#endif

#ifdef GPU_SPRNG

//perturb the particles
__kernel void perturb(__global float* x,
                      __global float* y,
                      __global float* z,
                      __global rngen* generators,
                      __global unsigned int* curBuffer,
                      unsigned int numWalkers,
                      unsigned int numParticles)
{
  unsigned int thisWalker = get_global_id(0);
  if(thisWalker < numWalkers)
  {
    unsigned int thisBuffer = curBuffer[thisWalker];
    unsigned int proposedBuffer = thisBuffer == 0 ? 1 : 0;

    __global float* curX = &x[thisBuffer * numWalkers * numParticles + 
      thisWalker * numParticles];
    __global float* curY = &y[thisBuffer * numWalkers * numParticles + 
      thisWalker * numParticles];
    __global float* curZ = &z[thisBuffer * numWalkers * numParticles + 
      thisWalker * numParticles];
    __global float* proposedX = &x[proposedBuffer * numWalkers * numParticles + 
      thisWalker * numParticles];
    __global float* proposedY = &y[proposedBuffer * numWalkers * numParticles + 
      thisWalker * numParticles];
    __global float* proposedZ = &z[proposedBuffer * numWalkers * numParticles + 
      thisWalker * numParticles];

    for(unsigned int curParticle = 0; curParticle < numParticles; curParticle++)
    {
      proposedX[curParticle] =
        curX[curParticle] +
        (get_rn_flt(&generators[thisWalker], 0) - 0.5) * stepSize;
      proposedY[curParticle] = 
        curY[curParticle] + 
        (get_rn_flt(&generators[thisWalker], 0) - 0.5) * stepSize;
      proposedZ[curParticle] = 
        curZ[curParticle] + 
        (get_rn_flt(&generators[thisWalker], 0) - 0.5) * stepSize;
    }
  }
}

#else

//perturb the particles (Using random numbers from CPU)
__kernel void perturb(__global float* x,
                      __global float* y,
                      __global float* z,
                      __global float* randomNumbers,
                      __global unsigned int* curBuffer,
                      unsigned int numWalkers,
                      unsigned int numParticles)
{
  unsigned int thisWalker = get_global_id(0);
  if(thisWalker < numWalkers)
  {
    unsigned int thisBuffer = curBuffer[thisWalker];
    unsigned int proposedBuffer = thisBuffer == 0 ? 1 : 0;

    __global float* curX = &x[thisBuffer * numWalkers * numParticles + 
      thisWalker * numParticles];
    __global float* curY = &y[thisBuffer * numWalkers * numParticles + 
      thisWalker * numParticles];
    __global float* curZ = &z[thisBuffer * numWalkers * numParticles + 
      thisWalker * numParticles];
    __global float* proposedX = &x[proposedBuffer * numWalkers * numParticles + 
      thisWalker * numParticles];
    __global float* proposedY = &y[proposedBuffer * numWalkers * numParticles + 
      thisWalker * numParticles];
    __global float* proposedZ = &z[proposedBuffer * numWalkers * numParticles + 
      thisWalker * numParticles];

    __global float* randomPtr = &randomNumbers[thisWalker];

    for(unsigned int curParticle = 0; curParticle < numParticles; curParticle++)
    {
      proposedX[curParticle] =
        curX[curParticle] +
        (randomPtr[0 * numWalkers] - 0.5) * stepSize;
      proposedY[curParticle] =
        curY[curParticle] +
        (randomPtr[1 * numWalkers] - 0.5) * stepSize;
      proposedZ[curParticle] =
        curZ[curParticle] +
        (randomPtr[2 * numWalkers] - 0.5) * stepSize;
      randomPtr += 3 * numWalkers;
    }
  }
}
#endif

//combine our results...
//***This function is anything but performance optimized***
__kernel void reducePartials(__global double* partialWaveFunction,
                             __global double* partialPotential,
                             __global double* partialKinetic,
                             __global double* qForceX,
                             __global double* qForceY,
                             __global double* qForceZ,
                             __global unsigned int* curBuffer,
                             unsigned int numWalkers,
                             unsigned int numParticles,
                             __global double* waveFunction,
                             __global double* potential,
                             __global double* kinetic)
{
  unsigned int thisWalker = get_global_id(0);
  double totalKinetic = 0.0;
  double totalPotential = 0.0;
  double totalWaveFunction = 0.0;
  double qForceXReg;
  double qForceYReg;
  double qForceZReg;

  partialWaveFunction = &partialWaveFunction[thisWalker * numParticles];
  partialPotential = &partialPotential[thisWalker * numParticles];
  partialKinetic = &partialKinetic[thisWalker * numParticles];
  qForceX = &qForceX[thisWalker * numParticles];
  qForceY = &qForceY[thisWalker * numParticles];
  qForceZ = &qForceZ[thisWalker * numParticles];
  
  if(thisWalker < numWalkers)
  {
    //Sum total energies and wavefunction
    for(unsigned int curParticle = 0; curParticle < numParticles; curParticle++)
    {
      totalWaveFunction += partialWaveFunction[curParticle];
      totalPotential += partialPotential[curParticle];
      totalKinetic += partialKinetic[curParticle];
    }

    totalKinetic *= 2.0;

    //Add in the qForces
    for(unsigned int curParticle = 0; curParticle < numParticles; curParticle++)
    {
      qForceXReg = qForceX[curParticle];
      qForceYReg = qForceY[curParticle];
      qForceZReg = qForceZ[curParticle];
      totalKinetic -= 
        qForceXReg * qForceXReg + 
        qForceYReg * qForceYReg +
        qForceZReg * qForceZReg;
    }

    totalKinetic *= 0.5 / PARTICLE_MASS;

    //Make the pointers correspond to this walker and current buffer
    unsigned int proposedBuffer = curBuffer[thisWalker] == 0 ? 1 : 0;
    unsigned int index = proposedBuffer * numWalkers + thisWalker;

    waveFunction[index] = totalWaveFunction;
    potential[index] = totalPotential;
    kinetic[index] = totalKinetic;
  }
}

#ifdef GPU_SPRNG
__kernel void accept(__global unsigned int* curBuffer,
                     unsigned int numWalkers,
                     __global rngen* generators,
                     __global double* waveFunction,
                     __global double* potential,
                     __global double* kinetic,
                     __global double* totalEnergy,
                     unsigned int shouldUpdateEnergy)
{
  unsigned int thisWalker = get_global_id(0);

  if(thisWalker < numWalkers)
  {
    unsigned int curBufferReg = curBuffer[thisWalker];
    unsigned int proposedBuffer = curBufferReg == 0 ? 1 : 0;
    double p;

    p = exp(2.0 * (waveFunction[proposedBuffer * numWalkers + thisWalker] -
                   waveFunction[curBufferReg * numWalkers + curBufferReg]));

    if(p >= 1.0 || get_rn_flt(&generators[thisWalker], 0) < p)
    {
      curBuffer[thisWalker] = proposedBuffer;

      totalEnergy[thisWalker] += 
        potential[proposedBuffer * numWalkers + thisWalker] + 
        kinetic[proposedBuffer * numWalkers + thisWalker];
    }
    else
    {
      totalEnergy[thisWalker] += 
        potential[curBufferReg * numWalkers + thisWalker] + 
        kinetic[curBufferReg * numWalkers + thisWalker];
    }
  }
}

#else

__kernel void accept(__global unsigned int* curBuffer,
                     unsigned int numWalkers,
                     __global float* randomNumbers,
                     __global double* waveFunction,
                     __global double* potential,
                     __global double* kinetic,
                     __global double* totalEnergy,
                     unsigned int shouldUpdateEnergy)
{
  unsigned int thisWalker = get_global_id(0);

  if(thisWalker < numWalkers)
  {
    unsigned int curBufferReg = curBuffer[thisWalker];
    unsigned int proposedBuffer = curBufferReg == 0 ? 1 : 0;
    double p;

    p = exp(2.0 * (waveFunction[proposedBuffer * numWalkers + thisWalker] -
                   waveFunction[curBufferReg * numWalkers + thisWalker]));

    if(p >= 1.0f || randomNumbers[thisWalker] < p)
    {
      curBuffer[thisWalker] = proposedBuffer;

      totalEnergy[thisWalker] += 
        potential[proposedBuffer * numWalkers + thisWalker] + 
        kinetic[proposedBuffer * numWalkers + thisWalker];
    }
    else
    {
      totalEnergy[thisWalker] +=
        potential[curBufferReg * numWalkers + thisWalker] + 
        kinetic[curBufferReg * numWalkers + thisWalker];
    }
  }
}

#endif

__kernel void computeEnergyPartials(__global float* x,
                                    __global float* y,
                                    __global float* z,
                                    unsigned int numParticles,
                                    __global unsigned int* curBuffer,
                                    __global double* waveFunction,
                                    __global double* potential,
                                    __global double* kinetic,
                                    __global double* qForceX,
                                    __global double* qForceY,
                                    __global double* qForceZ)
{
  unsigned int thisParticle = get_global_id(0);
  unsigned int thisWalker = get_global_id(1);
  unsigned int numWalkers = get_global_size(1);
  unsigned int otherParticle;
  unsigned int myWord = get_local_id(0);
  float dx;
  float dy;
  float dz;
  float thisX;
  float thisY;
  float thisZ;
  float otherX;
  float otherY;
  float otherZ;
  float rPow2;
  float r;
  float gradientDivR;

  //Accumulators
  double waveFunctionPartial = 0.0f;
  double potentialPartial = 0.0f;
  double thisQForceX = 0.0;
  double thisQForceY = 0.0;
  double thisQForceZ = 0.0;
  double kineticPartial = 0.0;
  
  //Local memory staging buffer to force coalescing
  __local float xTmp[64];
  __local float yTmp[64];
  __local float zTmp[64];

  //Get the proposed buffer we're reading from
  
  if(thisParticle < numParticles)
  {
    unsigned int proposedBuffer = curBuffer[thisWalker] == 0 ? 1 : 0;

    //Align the pointers to the correct walker and buffer
    x = &x[proposedBuffer * numParticles * numWalkers 
      + thisWalker * numParticles];
    y = &y[proposedBuffer * numParticles * numWalkers 
      + thisWalker * numParticles];
    z = &z[proposedBuffer * numParticles * numWalkers 
      + thisWalker * numParticles];
    
    waveFunction = &waveFunction[thisWalker * numParticles];
    potential = &potential[thisWalker * numParticles];
    kinetic = &kinetic[thisWalker * numParticles];
    qForceX = &qForceX[thisWalker * numParticles];
    qForceY = &qForceY[thisWalker * numParticles];
    qForceZ = &qForceZ[thisWalker * numParticles];

    //Fetch the coordinates of our particle
    thisX = x[thisParticle];
    thisY = y[thisParticle];
    thisZ = z[thisParticle];
    
    for(unsigned int curParticle = 0; 
        curParticle < numParticles; 
        curParticle += 64)
    {

      barrier(CLK_LOCAL_MEM_FENCE);

      if(curParticle + myWord < numParticles)
      {
        //Load x, y, z into shared memory. We do this to force coalesced reads
        //from global memory. Wait for loads to complete.
        //CurBlock << 5 <=> curBlock * 64
        xTmp[myWord] = x[curParticle + myWord];
        yTmp[myWord] = y[curParticle + myWord];
        zTmp[myWord] = z[curParticle + myWord];
      }

      barrier(CLK_LOCAL_MEM_FENCE);

      unsigned int lastElement = curParticle + 63 < numParticles ?
        64 :
        numParticles % 64;

      for(otherParticle = 0; otherParticle < lastElement; otherParticle++)
      {
        otherX = xTmp[otherParticle];
        otherY = yTmp[otherParticle];
        otherZ = zTmp[otherParticle];

        dx = thisX - otherX;
        dy = thisY - otherY;
        dz = thisZ - otherZ;

        rPow2 = dx * dx + dy * dy + dz * dz;
        r = sqrt(rPow2);

        gradientDivR = computeLogPsiPrime(r) / r;

        //Compute energies and wavefunction only if otherParticle < thisParticle
        if(curParticle + otherParticle < thisParticle)
        {
          waveFunctionPartial += (double)computeLogPsi(r);
          potentialPartial += (double)computePotential(r);
          kineticPartial -= 
            (double)(computeLogPsiDoublePrime(r) + 2.0f * gradientDivR);
        }

        //Compute qForces for all particles
        if(curParticle + otherParticle != thisParticle)
        {
          thisQForceX += (double)(dx * gradientDivR);
          thisQForceY += (double)(dy * gradientDivR);
          thisQForceZ += (double)(dz * gradientDivR);
        }
      }
    }

    waveFunction[thisParticle] = waveFunctionPartial;
    potential[thisParticle] = potentialPartial;
    qForceX[thisParticle] = thisQForceX;
    qForceY[thisParticle] = thisQForceY;
    qForceZ[thisParticle] = thisQForceZ;
    kinetic[thisParticle] = kineticPartial;
  }
}

